import json
import os.path
import platform
import shutil
import time

import requests
from PyQt6.QtCore import QThreadPool
from PyQt6.QtGui import QDragEnterEvent, QDropEvent
from PyQt6.QtWidgets import QPushButton, QWidget, QTextEdit, \
    QSizePolicy, QLabel, QScrollArea, QCheckBox, QLineEdit
from PyQt6.uic import loadUi

import global_config
import main_model
import midjourney_config
from global_config import *
from scripts import db_utils, beans
from scripts.WorkSingle import Worker
from scripts.configui.ImageGenerateTask import ImageTask, SVDTask, GenerateManager

from scripts.jianying import draft_parser
from ui.adapters.ScrollviewAdapter import ScrollviewAdapter
from ui.widgets.notice_dialog import ConfirmDialog


# 反推绘制页面（洗稿）
class WashPage(QWidget):
    def __init__(self):
        super().__init__()
        self.selectAll = False
        self.mv_dropview: QTextEdit = None
        self.srt_dropview: QTextEdit = None
        self.import_btn: QPushButton = None
        self.menu_btn: QPushButton = None
        self.back_btn: QPushButton = None
        self.save_btn: QPushButton = None
        self.loading: QLabel = None
        self.global_tv: QLabel = None
        self.global_cb: QLabel = None
        self.global_shuff_cb: QCheckBox = None
        self.parse_btn: QPushButton = None
        self.svd_btn: QPushButton = None
        self.jianying_btn: QPushButton = None
        self.generate_btn: QPushButton = None
        self.global_shuff_tv: QLineEdit = None
        self.storyboards: beans.Storyboard = None
        self.scrollview: QScrollArea = None
        self.scrollviewAdapter: ScrollviewAdapter = None
        self.videoDialog = None
        self.initView()
        self.imgGenManager = GenerateManager(self.on_generate_images)
        self.imgGenManager.startTask()

    def initView(self):
        self.setAcceptDrops(True)
        loadUi('./source/wash_page.ui', self)
        self.listAdapter = ScrollviewAdapter(self.scrollview)
        self.loading.hide()
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        # self.mv_dropview.setAcceptDrops(True)
        # self.srt_dropview.setAcceptDrops(True)

        self.video = get_video_path()
        self.srt = get_srt_path()
        if os.path.exists(self.video):
            self.mv_dropview.setText(self.video)
        if os.path.exists(self.srt):
            self.srt_dropview.setText(self.srt)

        self.import_btn.clicked.connect(self.import_videofile)
        self.menu_btn.clicked.connect(self.on_open_dir)
        self.back_btn.clicked.connect(self.on_pressback)
        self.save_btn.clicked.connect(self.update_workflow)
        self.parse_btn.clicked.connect(self.parse_movie2storyboard)
        self.generate_btn.clicked.connect(self.generate_checked_item)
        self.select_cb.clicked.connect(self.select_all_holder)
        self.svd_btn.clicked.connect(self.generate_svd)
        self.jianying_btn.clicked.connect(self.generate_jianying_draft)

        # 加载数据
        storyboards_path = get_storyboards_path()
        config_json_path = get_config_json()
        if os.path.exists(storyboards_path):
            self.storyboards = db_utils.read_workflow(storyboards_path)
            self.load_storyboards(self.storyboards)
        self.global_shuff_tv.setText(db_utils.read_config_json(config_json_path))

        # self.global_cb.clicked.connect()
        self.global_tv.setText(midjourney_config.positive_prompt)

    def select_all_holder(self):
        self.selectAll = False if self.selectAll else True
        for holder in self.listAdapter.itemholders:
            holder.setSelected(self.selectAll)

    # 加载解析完的 分镜数据
    def load_storyboards(self, storyboards):
        self.listAdapter.load(storyboards)
        # 创建一个信号，触发时传递当前时间给槽函数

    def parse_movie2storyboard(self):
        if not os.path.exists(get_srt_path()):
            ConfirmDialog('提示', '请导入原视频和字幕文件').get_result()
            return
        self.loading.show()

        def callback():
            self.storyboards = db_utils.read_workflow(get_storyboards_path())
            self.global_shuff_tv.setText(db_utils.read_config_json(get_config_json()))
            self.load_storyboards(self.storyboards)
            self.loading.hide()

        main_model.parse_movie2storyboard(callback)

    def on_pressback(self):
        self.parent().initUI()

    def on_open_dir(self):
        if platform.system() == 'Darwin':
            os.system(f'open /{workflows_dir}{get_cur_project()}')
        else:
            os.startfile(f'{workflows_dir}{get_cur_project()}')

    def import_videofile(self):
        self.video = get_video_path()
        self.srt = get_srt_path()
        msg = ''
        srt_path = self.srt_dropview.toPlainText()
        if srt_path == '':
            msg = '请输入字幕路径。'
        elif srt_path != self.srt:
            shutil.copy(srt_path, self.srt)
        else:
            msg = '字幕已存在。'
        mv_path = self.mv_dropview.toPlainText()
        if mv_path != '' and mv_path != self.video:
            shutil.copy(mv_path, self.video)
        elif mv_path != '':
            msg = msg + '视频已存在'
        if msg != '':
            dialog = ConfirmDialog('提示', msg)
            result = dialog.get_result()

    # 生成选中的分镜图片
    def generate_checked_item(self):
        def inner_generate():
            for holder in self.listAdapter.itemholders:
                if holder.index_checkbox.isChecked():
                    holder.setSelected(False)
                    if holder.bean.ref_img is None and holder.bean.local_img is not None:
                        print("upload")
                        res = requests.post(url=comfyui_upload,
                                            files={'image': open(holder.bean.local_img, 'rb')})
                        holder.bean.ref_img = res.json()['name']
                        db_utils.save_wordflow(get_storyboards_path(), self.storyboards)
                    qp_config = db_utils.load_project_config()
                    promots = holder.bean.to_full_prompts(self.global_tv.text(), self.global_shuff_tv.text())
                    # self.loading.show()
                    file = holder.bean.ref_img
                    # if not os.path.exists(file):
                    if file is not None:
                        self.imgGenManager.enQueue(ImageTask(promots, qp_config['negative_prompt'], holder.index, file,
                                                             width=qp_config['width'], height=qp_config['height']))
                    else:
                        self.imgGenManager.enQueue(ImageTask(promots, qp_config['negative_prompt'], holder.index,
                                                             width=qp_config['width'], height=qp_config['height']
                                                             ))
                    time.sleep(0.5)

        worker = Worker(inner=lambda: inner_generate(), callback=lambda a: ())
        thread = QThreadPool.globalInstance()
        thread.start(worker)

    # 图片转视频
    def generate_svd(self):
        def inner_generate():
            for holder in self.listAdapter.itemholders:
                if holder.index_checkbox.isChecked():
                    holder.setSelected(False)

                    if holder.bean.select_img is not None:
                        aiImgDir = f'{global_config.get_workflows_img_dir()}\\{holder.index}\\'
                        if os.path.exists(aiImgDir):
                            filelist = os.listdir(aiImgDir)
                            holder.bean.select_img = f'{aiImgDir}{filelist[0]}' if len(filelist) > 0 else None
                            db_utils.save_wordflow(get_storyboards_path(), self.storyboards)

                    print("upload")
                    res = requests.post(url=comfyui_upload,
                                        files={'image': open(holder.bean.select_img, 'rb')})
                    upload_img = res.json()['name']

                    self.imgGenManager.enQueue(SVDTask(holder.index, upload_img, holder.index))

                    time.sleep(0.1)

        worker = Worker(inner=lambda: inner_generate(), callback=lambda a: ())
        thread = QThreadPool.globalInstance()
        thread.start(worker)

    # 生成剪映草稿
    def generate_jianying_draft(self):
        draft_parser.sycn_draft()

    def update_workflow(self):
        storyboards_path = get_storyboards_path()
        db_utils.save_config_json(get_config_json(), self.global_shuff_tv.text())
        if self.storyboards is not None:
            db_utils.save_wordflow(storyboards_path, self.storyboards)

    def on_generate_images(self, success, result_images, taskdir, index, isEnd, isVideo):
        if isEnd:
            self.loading.hide()
        if not success: return

        print(index, result_images)
        for holder in self.listAdapter.itemholders:
            if holder.index == index:
                # 自动选中第一个图片
                if holder.bean.select_img is None or not os.path.exists(holder.bean.select_img):
                    aiImgDir = f'{global_config.get_workflows_img_dir()}\\{index}\\'
                    if os.path.exists(aiImgDir):
                        filelist = os.listdir(aiImgDir)
                        holder.bean.select_img = f'{aiImgDir}{filelist[0]}' if len(filelist) > 0 else None
                db_utils.save_wordflow(get_storyboards_path(), self.storyboards)
                holder.update_imagesUI()

        if isVideo:
            if self.videoDialog == None:
                self.videoDialog = ConfirmDialog('提示', f'分镜{index}，视频生成完毕，打开目录查看')
            self.videoDialog.message_label.setText(f'分镜{index}，视频生成完毕，打开目录查看')
            # self.videoDialog.get_result()

    def dragEnterEvent(self, e: QDragEnterEvent):
        e.acceptProposedAction()

    def dropEvent(self, e: QDropEvent):
        for url in e.mimeData().urls():
            file_name = url.toLocalFile()
            if 'srt' in file_name:
                self.srt_dropview.setText(file_name)
            elif 'mp4' in file_name:
                self.mv_dropview.setText(file_name)

    def resizeEvent(self, event):
        # 获取新的窗口尺寸
        size = self.size()
        print(f"窗口尺寸变化为宽度: {size.width()}， 高度: {size.height()}", event)
        geo = self.loading.geometry()
        geo.setWidth(size.width())
        geo.setHeight(size.height())
        self.loading.setGeometry(geo)
